#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelMishGrad
{
public:
    __aicore__ inline KernelMishGrad() {}
    __aicore__ inline void Init(GM_ADDR grad, GM_ADDR x, GM_ADDR tanhx, GM_ADDR y, bool tanhx_opt,
                                uint32_t totalNum, uint32_t tileNum, uint32_t tileLength, uint32_t lasttileLength)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->totalNum = totalNum;
        this->tileNum = tileNum;
        this->tileLength = tileLength;
        this->lasttileLength = lasttileLength;
        this->tanhx_opt = tanhx_opt;

        gGm.SetGlobalBuffer((__gm__ DTYPE_X *)grad, this->totalNum);
        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, this->totalNum);
        tGm.SetGlobalBuffer((__gm__ DTYPE_X *)tanhx, this->totalNum);
        yGm.SetGlobalBuffer((__gm__ DTYPE_X *)y, this->totalNum);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(inQueueG, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(inQueueT, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(CalcPow, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(CalcAdd, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(CalcT1, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(CalcT2, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Y));
    }
    __aicore__ inline void Process()
    {

        int32_t loopCount = this->tileNum;
        if (this->tanhx_opt)
        {
            for (int32_t i = 0; i < loopCount; i++)
            {
                if (i == this->tileNum - 1)
                {
                    this->processDataNum = this->lasttileLength;
                }

                CopyInwithT(i);
                ComputewithT(i);
                CopyOut(i);
            }
        }
        else
        {
            for (int32_t i = 0; i < loopCount; i++)
            {
                if (i == this->tileNum - 1)
                {
                    this->processDataNum = this->lasttileLength;
                }
                CopyIn(i);
                Compute(i);
                CopyOut(i);
            }
        }
    }

private:
    __aicore__ inline void CopyInwithT(int32_t progress)
    {

        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X> gLocal = inQueueG.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X> tLocal = inQueueT.AllocTensor<DTYPE_X>();

        DataCopy(xLocal, xGm[progress * this->tileLength], this->processDataNum);
        DataCopy(gLocal, gGm[progress * this->tileLength], this->processDataNum);
        DataCopy(tLocal, tGm[progress * this->tileLength], this->processDataNum);

        inQueueX.EnQue(xLocal);
        inQueueG.EnQue(gLocal);
        inQueueT.EnQue(tLocal);
    }
    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X> gLocal = inQueueG.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[progress * this->tileLength], this->processDataNum);
        DataCopy(gLocal, gGm[progress * this->tileLength], this->processDataNum);
        inQueueX.EnQue(xLocal);
        inQueueG.EnQue(gLocal);
    }
    __aicore__ inline void ComputewithT(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> gLocal = inQueueG.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> tLocal = inQueueT.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> powLocal = CalcPow.Get<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();

        // Compute
        Mul(powLocal, tLocal, tLocal, this->processDataNum);
        Muls(powLocal, powLocal, (DTYPE_X)-1, this->processDataNum);
        Adds(powLocal, powLocal, (DTYPE_X)1, this->processDataNum);
        Mul(powLocal, powLocal, xLocal, this->processDataNum);

        Muls(xLocal,xLocal,(DTYPE_X) -1, this->processDataNum);
        Exp(xLocal, xLocal, this->processDataNum);

        Adds(xLocal, xLocal, (DTYPE_X)1, this->processDataNum);
        Div(yLocal, powLocal, xLocal, this->processDataNum);
        Add(yLocal, yLocal, tLocal, this->processDataNum);
        Mul(yLocal, yLocal, gLocal, this->processDataNum);

        outQueueY.EnQue<DTYPE_Y>(yLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueG.FreeTensor(gLocal);
        inQueueT.FreeTensor(tLocal);
    }
    __aicore__ inline void Compute(int32_t progress){
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> gLocal = inQueueG.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> powLocal = CalcPow.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> addLocal = CalcAdd.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> Tmp1 = CalcT1.Get<DTYPE_X>();
        LocalTensor<DTYPE_X> Tmp2 = CalcT2.Get<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        
        Exp(powLocal,xLocal,this->processDataNum);
        Adds(addLocal,powLocal,(DTYPE_X) 1,this->processDataNum);

        Muls(Tmp1, xLocal, (DTYPE_X)4, this->processDataNum);
        Reciprocal(Tmp2, addLocal,this->processDataNum);
        Add(Tmp2,Tmp2,addLocal,this->processDataNum);
        Div(Tmp1,Tmp1,Tmp2,this->processDataNum);
        Mul(Tmp2,addLocal,addLocal,this->processDataNum);
        Adds(Tmp2,Tmp2,(DTYPE_X)1,this->processDataNum);
        Div(Tmp1,Tmp1,Tmp2,this->processDataNum);
        Muls(yLocal,yLocal,(DTYPE_X)0,this->processDataNum);
        Adds(yLocal,yLocal,(DTYPE_X)-2,this->processDataNum);
        Div(yLocal,yLocal,Tmp2,this->processDataNum);
        Add(yLocal,yLocal,Tmp1,this->processDataNum);
        Adds(yLocal,yLocal,(DTYPE_X)1,this->processDataNum);
        
        outQueueY.EnQue<DTYPE_Y>(yLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueG.FreeTensor(gLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(yGm[progress * this->tileLength], yLocal, this->processDataNum);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueG, inQueueT;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> CalcPow,CalcAdd, CalcT1, CalcT2;

    GlobalTensor<DTYPE_X> xGm, gGm, tGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t totalNum;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lasttileLength;
    uint32_t tanhx_opt;
    uint32_t processDataNum;
};

extern "C" __global__ __aicore__ void mish_grad(GM_ADDR grad, GM_ADDR x, GM_ADDR tanhx, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelMishGrad op;
    op.Init(grad, x, tanhx, y, tiling_data.tanhx_opt,
            tiling_data.totalNum, tiling_data.tileNum, tiling_data.tileLength, tiling_data.lasttileNum);
    op.Process();
}